from tkinter import N
from ..config import Base

from .BaseDb import BaseDb
from sqlalchemy import create_engine, Column, Integer, String, Date, DECIMAL, BIGINT, UniqueConstraint, Index, func
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from .model import StockDaily
import pandas as pd
from sqlalchemy.inspection import inspect
from typing import List



class DbHandler:
    def __init__(self, db: BaseDb):
        Base.metadata.create_all(db.engine)  # 创建所有表
        self.Session = sessionmaker(bind=db.engine)
        self.session = self.Session()
    
    def add_stock_daily(self, stock_daily):
        try:
            self.session.add(stock_daily)
            self.session.commit()
        except Exception as e:
            self.session.rollback()
            raise e
    
    def add_stock_daily_ignore_duplicate(self, stock_daily):
        """插入数据，如果已存在则忽略（不抛出异常）"""
        try:
            self.session.add(stock_daily)
            self.session.commit()
        except Exception as e:
            self.session.rollback()
            # 如果是唯一约束冲突，忽略；其他异常继续抛出
            if 'Duplicate entry' in str(e) or 'uk_stock_date' in str(e) or 'Unique' in str(e):
                # 数据已存在，忽略
                pass
            else:
                raise e
    
    def get_stock_daily(
            self, 
            ts_code: str = None, 
            start_date: str = None, 
            end_date: str = None
        ):
        query = self.session.query(StockDaily)
        
        if ts_code:
            query = query.filter(StockDaily.ts_code == ts_code).order_by(StockDaily.trade_date.asc())  # 时间从小到大
        
        if start_date and end_date:
            query = query.filter(StockDaily.trade_date.between(start_date, end_date)).order_by(StockDaily.trade_date.asc())  # 时间从小到大
        elif start_date:
            query = query.filter(StockDaily.trade_date >= start_date).order_by(StockDaily.trade_date.asc())  # 时间从小到大
        elif end_date:
            query = query.filter(StockDaily.trade_date <= end_date).order_by(StockDaily.trade_date.asc())  # 时间从小到大
        
        return query.all()
    
    def get_stock_daily_df(self, list_stock_daily):
        if list_stock_daily:
            columns_order = [c.key for c in inspect(StockDaily).columns if c.key != 'id' and c.key != 'ts_code']
            df = pd.DataFrame([r.__dict__ for r in list_stock_daily])
            # 删除 SQLAlchemy 内部属性 _sa_instance_state
            df.drop(columns=['_sa_instance_state'], inplace=True, errors='ignore')
            # df.drop(columns=['id'], inplace=True, errors='ignore')
            df = df.reindex(columns=columns_order)
            return df
        else:
            return pd.DataFrame()  # 空 DataFrame
    
    def get_all_ts_codes(self)->List[str]:
        result = self.session.query(StockDaily.ts_code).distinct().all()
        # result 是列表 of tuple，需要取第一个元素
        return [r[0] for r in result]
    
    def get_latest_trade_date(self, ts_code: str = None):
        query = self.session.query(func.max(StockDaily.trade_date))
        if ts_code:
            query = query.filter(StockDaily.ts_code == ts_code)
        return query.scalar()
    
    def get_last_stock_daily(self, ts_code: str):
        return (
            self.session.query(StockDaily)
            .filter(StockDaily.ts_code == ts_code)
            .order_by(StockDaily.trade_date.desc())
            .first()  # 取排序后的第一条
        )

    def get_mainboard_ts_codes(self):
        """
        获取沪深主板股票的所有 ts_code（去重）
        返回列表，例如 ['600519.SH', '000001.SZ', ...]
        """
        result = self.session.query(StockDaily.ts_code).filter(
            (StockDaily.ts_code.like('600%.SH')) |
            (StockDaily.ts_code.like('601%.SH')) |
            (StockDaily.ts_code.like('603%.SH')) |
            (StockDaily.ts_code.like('605%.SH')) |
            (StockDaily.ts_code.like('000%.SZ')) |
            (StockDaily.ts_code.like('001%.SZ')) |
            (StockDaily.ts_code.like('002%.SZ'))
        ).distinct().all()

        # SQLAlchemy 返回的是列表 of tuple，需要取第一个元素
        return [r[0] for r in result]
    

    
    # def get_stock_daily(self, ts_code):
    #     return self.session.query(StockDaily).filter(StockDaily.ts_code == ts_code)



    # row = StockDaily(ts_code="000001.SZ", trade_date="20250103", open=100, high=101, low=99, close=100, pre_close=99, change=1, pct_chg=0.01, vol=1000, amount=10000)
    # stock_daily.add_stock_daily(row)

    # rows = stock_daily.get_stock_daily("000001.SZ")
    # for row in rows:
    #     print(row.id, row.ts_code, row.trade_date, row.open, row.close)


    # stock_daily.add_stock_daily(stock_daily)
    # stock_daily.session.commit()
